#Copyright (C) Nial Peters 2012
#
#This file is part of AvoScan.
#
#AvoScan is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.
#
#AvoScan is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with AvoScan.  If not, see <http://www.gnu.org/licenses/>.


import pylab as plt
import numpy
import datetime
import cPickle
from avoscan import processing
import wx
import sys
from std_ops.iter_ import parallel_process
import scipy
import math
import os.path
from ica import cPlume
from ica.ica_calc import ProcessedScan, fullscans


class DraggableLine:
    """
    This class is mostly based on the draggable rectangles example from 
    http://matplotlib.sourceforge.net/users/event_handling.html
    
    It provides a vertical line across 'axis' which can be dragged to any x
    position. Useful for picking limits.
    
        * x_pos - initial x position to draw line at
        * axis - axis to draw line into
        * callback - list of functions to run when the line gets moved. The functions
                     will be run in order, and will be passed the new x position of 
                     the line as their only argument.
    """
    
    lock = None
    
    def __init__(self, x_pos, axis, linewidth=2.0, color='r', callback=None):
        
        self.line = axis.plot([x_pos]*2,axis.get_ylim(), linewidth=linewidth, color=color)[0]

        self.press = None
        self.background = None
        self.callback = callback
        self.__connect()

   
    def __connect(self):
        #connect to all the events we need
        self.cidpress = self.line.figure.canvas.mpl_connect('button_press_event', self.on_press)
        self.cidrelease = self.line.figure.canvas.mpl_connect('button_release_event', self.on_release)
        self.cidmotion = self.line.figure.canvas.mpl_connect('motion_notify_event', self.on_motion)           
        self.cidredraw = self.line.figure.canvas.mpl_connect('draw_event', self.on_redraw)


    def disconnect(self):
        #disconnect all the stored connection ids
        self.line.figure.canvas.mpl_disconnect(self.cidpress)
        self.line.figure.canvas.mpl_disconnect(self.cidrelease)
        self.line.figure.canvas.mpl_disconnect(self.cidmotion)
        self.line.figure.canvas.mpl_disconnect(self.cidredraw)

        
    def on_press(self, event):
        #on button press we will see if the mouse is over us and store some data
        if event.inaxes != self.line.axes: return
        if DraggableLine.lock is not None: return
        contains, attrd = self.line.contains(event)
        if not contains: return

        x0 = self.line.get_xdata()[0],
        self.press = x0, event.xdata, event.ydata
        DraggableLine.lock = self

        # draw everything but the selected rectangle and store the pixel buffer
        canvas = self.line.figure.canvas
        axes = self.line.axes
        self.line.set_animated(True)
        canvas.draw()
        self.background = canvas.copy_from_bbox(self.line.axes.bbox)

        # now redraw just the rectangle
        axes.draw_artist(self.line)

        # and blit just the redrawn area
        canvas.blit(axes.bbox)


    def get_xpos(self):
        """
        Returns the current x position of the line.
        """
        return self.line.get_xdata()[0]
    
    
    def set_xpos(self,x_pos):
        """
        Sets the current x position of the line and redraws it.
        """
        self.line.set_xdata([x_pos]*2)
        self.line.axes.redraw_in_frame()
        
    
    def on_redraw(self, event):
        #reset the line's ylims so that you can't zoom to outside the size of the line
        self.line.set_ydata(self.line.axes.get_ylim())
        
        canvas = self.line.figure.canvas
        axes = self.line.axes        
        
        # restore the background region
        canvas.restore_region(self.background)

        # redraw just the current rectangle
        axes.draw_artist(self.line)

        # blit just the redrawn area
        canvas.blit(axes.bbox)
    
    
    def on_motion(self, event):
        #on motion we will move the rect if the mouse is over us
        if DraggableLine.lock is not self:
            return
        if event.inaxes != self.line.axes: return
        x0, xpress, ypress = self.press
        dx = event.xdata - xpress
        dy = event.ydata - ypress

        self.line.set_xdata([x0[0] + dx]*2)
        self.line.set_ydata(self.line.axes.get_ylim())
        self.line.set_linestyle('--')

        canvas = self.line.figure.canvas
        axes = self.line.axes
        # restore the background region
        canvas.restore_region(self.background)

        # redraw just the current rectangle
        axes.draw_artist(self.line)

        # blit just the redrawn area
        canvas.blit(axes.bbox)

 
    def on_release(self, event):
        #on release we reset the press data
        if DraggableLine.lock is not self:
            return
        
        self.line.set_linestyle('-')
        self.press = None
        DraggableLine.lock = None

        # turn off the rect animation property and reset the background
        self.line.set_animated(False)
        self.background = None

        # redraw the full figure
        self.line.figure.canvas.draw()
        
        if self.callback is not None:
            self.callback(event.xdata)
 

class ScanSelector:
    def __init__(self, scans_list, ax, callback=None):
        self.ax = ax
        self.plotted_pts = {}
        self.callback = callback
        self.scans = {}
        for s in scans_list:
            self.scans[s.ica_time] = s
        
        self.times = numpy.array([s.ica_time for s in scans_list])
        self.times.sort()
        self.previous_selection = None
        
        self.update_plot()
        #start off with the first point selected
        self.select(self.scans[numpy.min(self.times)])

        ax.figure.canvas.mpl_connect('pick_event', self.onpick)
        #ax.set_title("Integrated Column Amount")
    
    
    def update_key_for_scan(self, scan, old_key):
        self.scans.pop(old_key)
        self.scans[scan.ica_time] = scan
        self.times[numpy.where(self.times == old_key)] = scan.ica_time
        v = self.plotted_pts.pop(old_key)
        self.plotted_pts[scan.ica_time] = v
        
        if self.previous_selection.ica_time == old_key:
            self.previous_selection = scan
    
    def update_plot(self):
        self.ax.clear()
        self.plotted_pts = {}
        for s in self.scans.values():
            self.plotted_pts[s.ica_time] = self.ax.errorbar(s.ica_time, s.ica, s.ica_err, s.ica_time_err, 'b.',picker=3)
        if self.previous_selection is not None:
            self.select(self.previous_selection, force_redraw=True)
        
    def delete_current_scan(self):
        self.scans.pop(self.previous_selection.ica_time)
        self.times = self.times[numpy.where(self.times != self.previous_selection.ica_time)]
        self.previous_selection = None
        self.update_plot()

        self.select(self.scans[numpy.min(self.times)])
        
    def get_current_selection(self):
        return self.previous_selection
    
    def get_current_scan_num(self):
        #returns a tuple (cur_scan_num, n) where n is the total number of scans
        n = len(self.times)
        cur_scan_num = numpy.where(self.times == self.previous_selection.ica_time)[0][0] + 1 #want to start counting at 1
        return (cur_scan_num, n)
        
    def onpick(self, event):
        x_pos = event.artist.get_xdata()[0]
        scan = self.scans[x_pos]
        self.select(scan)
        
    
    def select_next(self):
        current_time_idx = numpy.where(self.times == self.previous_selection.ica_time)[0][0]
        if current_time_idx + 1 == len(self.times):
            #go back to the beginning
            next_scan = self.scans[self.times[0]]
        else:
            next_scan = self.scans[self.times[current_time_idx + 1]]
        self.select(next_scan)
    
    def select_prev(self):
        current_time_idx = numpy.where(self.times == self.previous_selection.ica_time)[0][0]
        if current_time_idx == 0:
            #go to the end
            prev_scan = self.scans[self.times[-1]]
        else:
            prev_scan = self.scans[self.times[current_time_idx - 1]]
        self.select(prev_scan)
                
    def select(self, scan, force_redraw=False):
        if scan == self.previous_selection and not force_redraw:
            return
        
        plt_pt = self.plotted_pts[scan.ica_time]
        plt_pt[0].set_color('r')
        for p in plt_pt[1]:
            p.set_color('r')
        for p in plt_pt[2]:
            p.set_color(((1,0,0),))
               
        #change old selection back to blue
        if self.previous_selection is not None and scan != self.previous_selection:
            old_pt = self.plotted_pts[self.previous_selection.ica_time]
            old_pt[0].set_color('b')
            for p in old_pt[1]:
                p.set_color('b')
            for p in old_pt[2]:
                p.set_color(((0,0,1),))
        
        plt_pt[0].figure.canvas.draw()
        self.previous_selection = scan
        
        #run the callback function
        if self.callback is not None:
            self.callback(scan)


class BackgroundRangeSelector:
    def __init__(self, ax):
        self.ax = ax
        self.lim_line1 = None
        self.lim_line2 = None
        self.plot = []
    
    def update(self, scan):
        if self.plot:
            self.plot[0].remove()
            for p in self.plot[1]:
                p.remove()
            for p in self.plot[2]:
                p.remove()
        if self.lim_line1 is None:
                self.lim_line1 = DraggableLine(scan.bkgd_lowlim, self.ax, color='g')
                self.lim_line2 = DraggableLine(scan.bkgd_highlim, self.ax, color='g')
        else:
            self.lim_line1.set_xpos(scan.bkgd_lowlim)
            self.lim_line2.set_xpos(scan.bkgd_highlim)

        self.plot = self.ax.errorbar(scan.angles, scan.col_amount, scan.col_err, scan.angle_err,'b+')
        
    def get_lims(self):
        pos1 = self.lim_line1.get_xpos()
        pos2 = self.lim_line2.get_xpos()
        
        return min(pos1,pos2), max(pos1,pos2)



class BackgroundFitDisplay:
    def __init__(self, ax):
        self.ax = ax
    
    def update(self, scan):
        self.ax.clear()
        scan.plot_bkgd_fit(self.ax)



class RadiusFitDisplay:
    def __init__(self, ax):
        self.ax = ax
    
    def update(self, scan):
        self.ax.clear()
        scan.plot_radius_fit(self.ax)        



class PlotManager:
    def __init__(self,scans):
        self.callbacks = []
        self.int_method = scans[0].integration_method
        
        #create the plotting window
        bkgd_select_ax = plt.subplot2grid((3,2), (0,0), colspan=2)
        bkgd_fit_ax = plt.subplot2grid((3,2), (1,0))
        
        ica_ax = plt.subplot2grid((3,2), (2, 0), colspan=2)
        self.plots = []
        
        self.bkgd_selector = BackgroundRangeSelector(bkgd_select_ax)
        self.plots.append(self.bkgd_selector)
        
        self.bkgd_fit = BackgroundFitDisplay(bkgd_fit_ax)
        self.plots.append(self.bkgd_fit)
        
        if (self.int_method == "Cylinder Fit"):
            rad_est_ax = plt.subplot2grid((3,2), (1, 1))
            self.rad_fit = RadiusFitDisplay(rad_est_ax)
            self.plots.append(self.rad_fit)
        
        
        self.scan_selector = ScanSelector(scans, ica_ax, callback=self.update)
        #self.plots.append(self.scan_selector)
        
        
    
    def delete_current_scan(self):
        self.scan_selector.delete_current_scan()    
        
    def register_callback(self, f):
        self.callbacks.append(f)
        
    def get_scans(self):
        return self.scan_selector.scans.values()
    
    def next(self):
        self.scan_selector.select_next()
        
    def prev(self):
        self.scan_selector.select_prev()
    
    def show(self):    
        plt.show()
    
    def get_current_scan_num(self):
        #returns a tuple (cur_selection, n) where n is the total number of scans
        return self.scan_selector.get_current_scan_num() 
     
    def get_current_scan(self):
        return self.scan_selector.get_current_selection()
             
    def update(self, scan, old_key=None):
        #if the updated scan has a new ica_time, then we need to update the dict
        if old_key is not None:
            self.scan_selector.update_key_for_scan(scan, old_key)
        
        wx.BeginBusyCursor()
        for p in self.plots:
            p.update(scan)
        
        plt.draw()
        
        for f in self.callbacks:
            f(scan)

        wx.EndBusyCursor()


class ControlWindow(wx.Frame):
    def __init__(self, plot_manager, title):
        self.plot_manager = plot_manager
        wx.Frame.__init__(self, None, wx.ID_ANY, title)
        self.top_panel = wx.Panel(self, wx.ID_ANY)
        
        self.main_hsizer = wx.BoxSizer(wx.HORIZONTAL)
        self.main_sizer = wx.BoxSizer(wx.VERTICAL)
        
        self.main_sizer.Add(SkipControls(self.top_panel, self.plot_manager),0,wx.ALIGN_TOP|wx.ALIGN_CENTER_HORIZONTAL)
        
        self.main_sizer.Add(wx.StaticText(self.top_panel,wx.ID_ANY, "Current Scan:"))
        self.main_sizer.Add(CurrentScanValues(self.top_panel, self.plot_manager),0,wx.EXPAND)
        
        self.main_sizer.AddSpacer(10)
        self.main_sizer.Add(wx.StaticText(self.top_panel,wx.ID_ANY, "Background Fitting:"))
        self.main_sizer.Add(BackgroundFittingControls(self.top_panel, self.plot_manager),0,wx.EXPAND)
        
        
        if plot_manager.int_method == "Cylinder Fit":
            self.main_sizer.AddSpacer(10)
            self.main_sizer.Add(wx.StaticText(self.top_panel,wx.ID_ANY, "Error Analysis:"))
            self.main_sizer.Add(ErrorCalculationControls(self.top_panel, self.plot_manager),0,wx.EXPAND)
        
        self.hsizer = wx.BoxSizer(wx.HORIZONTAL)
        
        self.save_all_button = wx.Button(self.top_panel, wx.ID_ANY, "Save All")
        self.delete_button = wx.Button(self.top_panel, wx.ID_ANY, "Delete Scan")
        self.export_button = wx.Button(self.top_panel, wx.ID_ANY, "Export All")
        self.main_sizer.AddSpacer(15)
        self.hsizer.Add(self.delete_button,0,wx.ALIGN_LEFT|wx.ALIGN_BOTTOM)
        self.hsizer.Add(self.export_button,0,wx.ALIGN_LEFT|wx.ALIGN_BOTTOM)
        self.hsizer.Add(self.save_all_button,0,wx.ALIGN_RIGHT|wx.ALIGN_BOTTOM)
        self.main_sizer.Add(self.hsizer,1,wx.ALIGN_BOTTOM|wx.ALIGN_RIGHT)
        self.main_sizer.AddSpacer(5)
        
        self.main_hsizer.AddSpacer(5)
        self.main_hsizer.Add(self.main_sizer, 1, wx.EXPAND)
        self.main_hsizer.AddSpacer(5)
        
        self.top_panel.SetSizer(self.main_hsizer)
        self.main_hsizer.Fit(self)
        self.top_panel.SetAutoLayout(1)
        
        wx.EVT_BUTTON(self, self.save_all_button.GetId(), self.on_save)
        wx.EVT_BUTTON(self, self.delete_button.GetId(), self.on_delete)
        wx.EVT_BUTTON(self, self.export_button.GetId(), self.on_export)
        wx.EVT_CLOSE(self, self.on_close)
        
        fig = plt.gcf()
        fig.canvas.mpl_connect('close_event', self.on_fig_close)

        
        self.Show()
    
    def on_save(self,evt):
        filename = wx.SaveFileSelector("Save","scans")
        if not filename:
            return
        wx.BeginBusyCursor()
        if not filename.endswith(".scans"):
            filename = filename + ".scans"
        with open(filename, "wb") as ofp:
            cPickle.dump(self.plot_manager.get_scans(), ofp)
        wx.EndBusyCursor()
    
    def on_export(self, evt):
        filename = wx.SaveFileSelector("Export as text file","txt")
        if not filename:
            return
        wx.BeginBusyCursor()
        if not filename.endswith(".txt"):
            filename = filename + ".txt"
        with open(filename, "w") as ofp:    
            ofp.write("# Date\t\tTime\t\tICA\t\tTime Error\t\tICA Error\n")
            for t in self.plot_manager.scan_selector.times:
                s = self.plot_manager.scan_selector.scans[t]
                ofp.write(str(s.ica_time)+"\t"+str(s.ica)+"\t"+str(s.ica_time_err)+"\t"+str(s.ica_err)+"\n")
        wx.EndBusyCursor()
                
          
          
    def on_delete(self,evt):
        self.plot_manager.delete_current_scan()
    
    def on_close(self,evnt):
        plt.close()
        self.Destroy()
        
    def on_fig_close(self,evnt):
        self.Destroy()
        
class CurrentScanValues(wx.Panel):
    def __init__(self, parent, plot_manager):
        self.plot_manager = plot_manager
        wx.Panel.__init__(self, parent, wx.ID_ANY, style=wx.BORDER_SIMPLE)
        
        self.hsizer = wx.BoxSizer(wx.HORIZONTAL)
        self.vsizer = wx.BoxSizer(wx.VERTICAL)
        
        self.hsizer.AddSpacer(10)
        self.hsizer.Add(self.vsizer,1, wx.EXPAND)
        
        self.ica_txt = wx.StaticText(self, wx.ID_ANY,"")
        self.time_txt = wx.StaticText(self, wx.ID_ANY,"")
        
        self.vsizer.AddSpacer(5)
        self.vsizer.Add(self.ica_txt,0,wx.ALIGN_LEFT)
        self.vsizer.AddSpacer(5)
        self.vsizer.Add(self.time_txt,0,wx.ALIGN_LEFT)
        self.vsizer.AddSpacer(5)
        
        self.SetSizer(self.hsizer)
        self.hsizer.Fit(self)
        self.SetAutoLayout(1)
        self.update(self.plot_manager.get_current_scan())
        self.plot_manager.register_callback(self.update)
        
    def update(self, scan):
        
        self.ica_txt.SetLabel("ICA = " + str(round(scan.ica,2)) + " +/- " + str(round(scan.ica_err,2)))
        
        self.time_txt.SetLabel("Time = "+str(scan.ica_time)+" +/- "+ str(scan.ica_time_err))

        


class  SkipControls(wx.Panel):
    def __init__(self, parent, plot_manager):
        wx.Panel.__init__(self, parent, wx.ID_ANY)
        self.plot_manager = plot_manager
        self.vsizer = wx.BoxSizer(wx.VERTICAL)
        self.hsizer = wx.BoxSizer(wx.HORIZONTAL)

        s,n = self.plot_manager.get_current_scan_num()
        self.txt = wx.StaticText(self,wx.ID_ANY,"Displaying scan "+str(s)+" of "+str(n))
        
        self.vsizer.AddSpacer(10)
        self.vsizer.Add(self.txt,0,wx.ALIGN_CENTER_HORIZONTAL)
        
        self.next_button = wx.Button(self, wx.ID_ANY, ">")
        self.prev_button = wx.Button(self, wx.ID_ANY, "<")
        
        self.hsizer.Add(self.prev_button,0,wx.ALIGN_LEFT)
        self.hsizer.Add(self.next_button,0,wx.ALIGN_RIGHT)
        
        self.vsizer.AddSpacer(5)        
        self.vsizer.Add(self.hsizer, 1, wx.ALIGN_CENTER_HORIZONTAL|wx.EXPAND)
        self.vsizer.AddSpacer(10)    
        wx.EVT_BUTTON(self, self.next_button.GetId(), self.on_next)
        wx.EVT_BUTTON(self, self.prev_button.GetId(), self.on_prev)
        
        self.plot_manager.register_callback(self.update)
        
        self.SetSizer(self.vsizer)
        self.vsizer.Fit(self)
        self.SetAutoLayout(1)
    
    def on_next(self,evt):
        self.next_button.Disable()
        self.prev_button.Disable()
        wx.Yield()
        self.plot_manager.next()
        wx.Yield()
        self.next_button.Enable()
        self.prev_button.Enable()

    
    def on_prev(self,evt):
        self.next_button.Disable()
        self.prev_button.Disable()
        wx.Yield()
        self.plot_manager.prev()
        wx.Yield()
        self.next_button.Enable()
        self.prev_button.Enable()

            
    def update(self, scan):
        s,n = self.plot_manager.get_current_scan_num()
        self.txt.SetLabel("Displaying scan "+str(s)+" of "+str(n))
        

class BackgroundFittingControls(wx.Panel):
    def __init__(self, parent, plot_manager):
        self.plot_manager = plot_manager
        wx.Panel.__init__(self, parent, wx.ID_ANY, style=wx.BORDER_SIMPLE)
        
        self.vsizer = wx.BoxSizer(wx.VERTICAL)
        self.autofit_check_box = wx.CheckBox(self, wx.ID_ANY, "Auto fit")
        
        self.vsizer.Add(self.autofit_check_box,0,wx.ALIGN_LEFT|wx.ALIGN_TOP)
        
        
        self.lims_hsizer = wx.BoxSizer(wx.HORIZONTAL)
        self.lower_lim_txtbox = wx.TextCtrl(self, wx.ID_ANY)
        self.upper_lim_txtbox = wx.TextCtrl(self, wx.ID_ANY)
        self.get_lims_button = wx.Button(self, wx.ID_ANY, "Get Limits")
        
        
        self.lims_hsizer.AddSpacer(5)
        self.lims_hsizer.Add(wx.StaticText(self, wx.ID_ANY,"Lower Limit:"),0, wx.ALIGN_CENTER_VERTICAL)
        self.lims_hsizer.Add(self.lower_lim_txtbox,0, wx.ALIGN_CENTER_VERTICAL)
        self.lims_hsizer.AddSpacer(10)
        self.lims_hsizer.Add(wx.StaticText(self, wx.ID_ANY,"Upper Limit:"),0, wx.ALIGN_CENTER_VERTICAL)
        self.lims_hsizer.Add(self.upper_lim_txtbox,0, wx.ALIGN_CENTER_VERTICAL)
        self.lims_hsizer.AddSpacer(10)
        self.lims_hsizer.Add(self.get_lims_button,0, wx.ALIGN_CENTER_VERTICAL)
        self.lims_hsizer.AddSpacer(10)
        self.vsizer.Add(self.lims_hsizer,1,wx.ALIGN_CENTER_HORIZONTAL)
        
        
        self.fit_order_sizer = wx.BoxSizer(wx.HORIZONTAL)
        self.fit_order_sizer.AddSpacer(5)
        self.fit_order_txtbox = wx.TextCtrl(self, wx.ID_ANY)
        
        self.fit_order_sizer.Add(wx.StaticText(self,wx.ID_ANY, "Fit Order:"),0, wx.ALIGN_CENTER_VERTICAL)
        self.fit_order_sizer.Add(self.fit_order_txtbox,0, wx.ALIGN_CENTER_VERTICAL)
        
        self.vsizer.Add(self.fit_order_sizer,1,wx.ALIGN_LEFT)
        
        self.apply_button = wx.Button(self, wx.ID_ANY, "Apply")
        self.vsizer.Add(self.apply_button,0,wx.ALIGN_BOTTOM|wx.ALIGN_RIGHT)
        
        
        wx.EVT_BUTTON(self, self.get_lims_button.GetId(), self.on_get_lims)
        wx.EVT_BUTTON(self, self.apply_button.GetId(), self.on_apply)
        wx.EVT_CHECKBOX(self, self.autofit_check_box.GetId(), self.on_autofit)
        self.SetSizer(self.vsizer)
        self.vsizer.Fit(self)
        self.SetAutoLayout(1)
        self.update(self.plot_manager.get_current_scan())
        self.plot_manager.register_callback(self.update)
        
    def on_apply(self,evt):
        wx.BeginBusyCursor()
        scan = self.plot_manager.get_current_scan()
        old_key = scan.ica_time
        scan.manual_bkgd = True
        scan.bkgd_lowlim = float(self.lower_lim_txtbox.GetValue())
        scan.bkgd_highlim = float(self.upper_lim_txtbox.GetValue())
        scan.bkgd_fit_order = int(self.fit_order_txtbox.GetValue())
        scan.calculate()
        print "ica = ",scan.ica
        self.plot_manager.update(scan, old_key = old_key)
        self.plot_manager.scan_selector.update_plot()
        wx.EndBusyCursor()
    
    def update(self,scan):        
        
        self.fit_order_txtbox.ChangeValue(str(scan.bkgd_fit_order))
        if not scan.manual_bkgd:
            self.autofit_check_box.SetValue(True)
            self.on_get_lims(None)
            self.set_non_editable()
        else:
            self.autofit_check_box.SetValue(False)
            self.lower_lim_txtbox.ChangeValue(str(round(scan.bkgd_lowlim,2)))
            self.upper_lim_txtbox.ChangeValue(str(round(scan.bkgd_highlim,2)))
            self.set_editable()
    
    def on_autofit(self,evt):
        if self.autofit_check_box.IsChecked():
            self.set_non_editable()
            wx.BeginBusyCursor()
            scan = self.plot_manager.get_current_scan()
            old_key = scan.ica_time
            scan.manual_bkgd = False
            scan.bkgd_fit_order = 1
            scan.calculate()
            self.plot_manager.update(scan, old_key = old_key)
            self.plot_manager.scan_selector.update_plot()
            wx.EndBusyCursor()
        else:
            self.set_editable()
    
    def set_non_editable(self):
        self.lower_lim_txtbox.Disable()
        self.upper_lim_txtbox.Disable()
        self.fit_order_txtbox.Disable()
        self.get_lims_button.Disable()
        self.apply_button.Disable()
    
    def set_editable(self):
        self.lower_lim_txtbox.Enable()
        self.upper_lim_txtbox.Enable()
        self.fit_order_txtbox.Enable()
        self.get_lims_button.Enable()
        self.apply_button.Enable()
        
    def on_get_lims(self, evt):
        
        lower_lim, upper_lim = self.plot_manager.bkgd_selector.get_lims()
        
        self.lower_lim_txtbox.ChangeValue(str(round(lower_lim,2)))
        self.upper_lim_txtbox.ChangeValue(str(round(upper_lim,2)))


class ErrorCalculationControls(wx.Panel):
    def __init__(self, parent, plot_manager):
        
        self.plot_manager = plot_manager
        wx.Panel.__init__(self, parent, wx.ID_ANY, style=wx.BORDER_SIMPLE)
        
        self.vsizer = wx.BoxSizer(wx.VERTICAL)
        self.error_check_box = wx.CheckBox(self, wx.ID_ANY, "None")
        
        self.vsizer.Add(self.error_check_box,0,wx.ALIGN_LEFT|wx.ALIGN_TOP)
                
        self.iter_hsizer = wx.BoxSizer(wx.HORIZONTAL)
        self.iter_txtbox = wx.TextCtrl(self, wx.ID_ANY)
       
        self.iter_hsizer.AddSpacer(5)
        self.iter_hsizer.Add(wx.StaticText(self, wx.ID_ANY,"Number of iterations:"),0, wx.ALIGN_CENTER_VERTICAL)
        self.iter_hsizer.Add(self.iter_txtbox,0, wx.ALIGN_CENTER_VERTICAL)
        self.iter_hsizer.AddSpacer(10)

        self.vsizer.Add(self.iter_hsizer)
        
        self.apply_button = wx.Button(self, wx.ID_ANY, "Apply")
        self.apply_all_button = wx.Button(self, wx.ID_ANY, "Apply To All")
        self.vsizer.Add(self.apply_button,0,wx.ALIGN_BOTTOM|wx.ALIGN_RIGHT)
        self.vsizer.Add(self.apply_all_button,0,wx.ALIGN_BOTTOM|wx.ALIGN_RIGHT)
        
        wx.EVT_CHECKBOX(self, self.error_check_box.GetId(), self.on_checkbox)
        wx.EVT_BUTTON(self, self.apply_button.GetId(), self.on_apply)
        wx.EVT_BUTTON(self, self.apply_all_button.GetId(), self.on_apply_all)
        
        self.SetSizer(self.vsizer)
        self.vsizer.Fit(self)
        self.SetAutoLayout(1)
        
        self.update(self.plot_manager.get_current_scan())
        self.plot_manager.register_callback(self.update)

        
    def update(self, scan):
        if scan.error_calc_iters == 1:
            self.error_check_box.SetValue(True)
            self.set_no_err()
            
        else:
            self.error_check_box.SetValue(False)
            self.set_err(scan)
    
    def set_no_err(self):
        self.iter_txtbox.SetValue("1")
        self.iter_txtbox.Disable()
        self.apply_button.Disable()
    
    def set_err(self, scan):
        self.iter_txtbox.SetValue(str(scan.error_calc_iters))
        self.iter_txtbox.Enable()
        self.apply_button.Enable()
        
        
    def on_checkbox(self, evnt):
        if self.error_check_box.IsChecked():
            self.set_no_err()
            wx.BeginBusyCursor()
            scan = self.plot_manager.get_current_scan()
            old_key = scan.ica_time
            scan.error_calc_iters = 1
            scan.calculate()
            self.plot_manager.update(scan, old_key = old_key)
            self.plot_manager.scan_selector.update_plot()
            wx.EndBusyCursor()
        else:
            self.set_err(self.plot_manager.get_current_scan())

    def on_apply(self, evnt):
        self.apply_button.Disable()
        self.apply_all_button.Disable()
        wx.Yield()
        wx.BeginBusyCursor()
        scan = self.plot_manager.get_current_scan()
        old_key = scan.ica_time
        scan.error_calc_iters = int(self.iter_txtbox.GetValue())
        scan.calculate()
        self.plot_manager.update(scan, old_key=old_key)
        self.plot_manager.scan_selector.update_plot()
        wx.EndBusyCursor()
        self.apply_button.Enable()
        self.apply_all_button.Enable()

    
    def on_apply_all(self, evnt):
        self.apply_button.Disable()
        self.apply_all_button.Disable()
        wx.Yield()
        wx.BeginBusyCursor()
        n_iters = int(self.iter_txtbox.GetValue())
        scans = self.plot_manager.scan_selector.scans.values()
        result = parallel_process(_calc_error, scans, n_iters)
        failed = []
        for scan,old_key in result:
            self.plot_manager.scan_selector.update_key_for_scan(scan, old_key)

            if scan.error_calc_iters != n_iters:
                failed.append(scan)
        self.plot_manager.scan_selector.update_plot()
        wx.EndBusyCursor()
        if failed:
            wx.MessageBox("Could not calculate errors for all scans. This is probably because there are not enough in-plume points to do the fitting. Try changing the background limits.", "ICA Inspect", wx.ICON_WARNING)
        self.apply_button.Enable()
        self.apply_all_button.Enable()
        
            

def _calc_error(scan, n_iters):
    old_key = scan.ica_time
    scan.error_calc_iters = n_iters
    scan.calculate()
    return (scan,old_key)
    
class FileChooser(wx.Frame):
    def __init__(self, filename):
        self.int_methods = ["Trapezoid Rule", "Cylinder Fit"]
        
        wx.Frame.__init__(self, None, wx.ID_ANY,"Open file - ICA Inspect")
        self.top_panel = wx.Panel(self, wx.ID_ANY)
        self.main_sizer = wx.BoxSizer(wx.VERTICAL)
        self.main_hsizer = wx.BoxSizer(wx.HORIZONTAL)
        self.main_hsizer.AddSpacer(5)
        self.main_hsizer.Add(self.main_sizer,1,wx.EXPAND)
        self.main_hsizer.AddSpacer(5)
                
        self.main_sizer.AddSpacer(5)
        
        self.filename_box = wx.TextCtrl(self.top_panel, wx.ID_ANY)
        self.filename_box.SetMinSize((350,self.filename_box.GetSize()[1]))
        self.browse_button = wx.Button(self.top_panel, wx.ID_ANY, "Browse")
        self.file_hsizer = wx.BoxSizer(wx.HORIZONTAL)
        
        self.file_hsizer.Add(self.filename_box,1,wx.ALIGN_CENTER_VERTICAL)
        self.file_hsizer.AddSpacer(5)
        self.file_hsizer.Add(self.browse_button,0,wx.ALIGN_RIGHT)   
        
        self.main_sizer.AddSpacer(10)
        self.main_sizer.Add(wx.StaticText(self.top_panel, wx.ID_ANY, "Choose a \'wangle\' or \'.scans\' file to open."),0,wx.ALIGN_CENTER_HORIZONTAL)
        self.main_sizer.AddSpacer(5)
        self.main_sizer.Add(wx.StaticText(self.top_panel,wx.ID_ANY,"Filename:"))
        self.main_sizer.Add(self.file_hsizer,0,wx.ALIGN_CENTER_HORIZONTAL|wx.EXPAND)
        self.main_sizer.AddSpacer(20)
        
        self.hdist_txtbox = wx.TextCtrl(self.top_panel, wx.ID_ANY)
        self.hdist_err_txtbox = wx.TextCtrl(self.top_panel, wx.ID_ANY)
        self.hdist_sizer = wx.BoxSizer(wx.HORIZONTAL)
        self.hdist_sizer.Add(wx.StaticText(self.top_panel,wx.ID_ANY,"Horizontal Distance:\t"),0, wx.ALIGN_CENTRE_VERTICAL)
        self.hdist_sizer.Add(self.hdist_txtbox,0, wx.ALIGN_CENTRE_VERTICAL)
        self.hdist_sizer.Add(wx.StaticText(self.top_panel,wx.ID_ANY,"\t+/-\t"),0, wx.ALIGN_CENTRE_VERTICAL)
        self.hdist_sizer.Add(self.hdist_err_txtbox,0, wx.ALIGN_CENTRE_VERTICAL)
        self.main_sizer.Add(self.hdist_sizer)
        self.main_sizer.AddSpacer(5)
 
        self.vangle_txtbox = wx.TextCtrl(self.top_panel, wx.ID_ANY)
        self.vangle_err_txtbox = wx.TextCtrl(self.top_panel, wx.ID_ANY)
        self.vangle_sizer = wx.BoxSizer(wx.HORIZONTAL)
        self.vangle_sizer.Add(wx.StaticText(self.top_panel,wx.ID_ANY,"Vertical Angle:\t\t"),0, wx.ALIGN_CENTRE_VERTICAL)
        self.vangle_sizer.Add(self.vangle_txtbox,0, wx.ALIGN_CENTRE_VERTICAL)
        self.vangle_sizer.Add(wx.StaticText(self.top_panel,wx.ID_ANY,"\t+/-\t"),0, wx.ALIGN_CENTRE_VERTICAL)
        self.vangle_sizer.Add(self.vangle_err_txtbox,0, wx.ALIGN_CENTRE_VERTICAL)
        self.main_sizer.Add(self.vangle_sizer)
        self.main_sizer.AddSpacer(5)
        
        self.conv_txtbox = wx.TextCtrl(self.top_panel, wx.ID_ANY)
        self.conv_sizer = wx.BoxSizer(wx.HORIZONTAL)
        self.conv_sizer.Add(wx.StaticText(self.top_panel,wx.ID_ANY,"Conversion Factor:\t"),0, wx.ALIGN_CENTRE_VERTICAL)
        self.conv_sizer.Add(self.conv_txtbox,0, wx.ALIGN_CENTRE_VERTICAL)
        self.main_sizer.Add(self.conv_sizer)
        self.main_sizer.AddSpacer(10)
        
        self.method_choice = wx.Choice(self.top_panel, wx.ID_ANY, choices=self.int_methods)
        self.method_sizer = wx.BoxSizer(wx.HORIZONTAL)
        self.method_sizer.Add(wx.StaticText(self.top_panel,wx.ID_ANY,"Integration Method:\t"),0, wx.ALIGN_CENTRE_VERTICAL)
        self.method_sizer.Add(self.method_choice,0, wx.ALIGN_CENTRE_VERTICAL)
        self.main_sizer.Add(self.method_sizer)
        self.main_sizer.AddSpacer(10)
           
        self.hdist_txtbox.Disable()
        self.hdist_err_txtbox.Disable()
        self.vangle_txtbox.Disable()
        self.vangle_err_txtbox.Disable()
        self.conv_txtbox.Disable()
        self.method_choice.Disable()
        
        self.main_sizer.AddStretchSpacer()
        self.ok_button = wx.Button(self.top_panel, wx.ID_ANY, "Ok")
        self.main_sizer.Add(self.ok_button,0,wx.ALIGN_RIGHT|wx.ALIGN_BOTTOM)  
        self.main_sizer.AddSpacer(5)
        
        self.top_panel.SetSizer(self.main_hsizer)
        self.main_hsizer.Fit(self)
        self.top_panel.SetAutoLayout(1)
        
        wx.EVT_BUTTON(self, self.browse_button.GetId(), self.on_browse)
        wx.EVT_BUTTON(self, self.ok_button.GetId(), self.on_ok)
        wx.EVT_TEXT(self, self.filename_box.GetId(), self.on_path_change)
        
        if filename is not None:
            self.filename_box.SetValue(filename)
            
            
        
    def on_browse(self, evnt):
        default_dir = None
        try:
            with open(os.path.normpath(os.path.expanduser('~/.ica_inspect_dir')),'r') as ifp:
                default_dir = ifp.read()
        
        except:
            pass
        
        if default_dir is not None and os.path.isdir(default_dir):
            dialog = wx.FileDialog(self, "Choose input file",default_dir)
        else:
            dialog = wx.FileDialog(self, "Choose input file")
        
        if dialog.ShowModal() != wx.CANCEL:
            self.filename_box.SetValue(dialog.GetPath())
            try:
                with open(os.path.normpath(os.path.expanduser('~/.ica_inspect_dir')),'w') as ofp:
                    ofp.write(os.path.dirname(dialog.GetPath()))
            except:
                pass
    
    def on_path_change(self, evnt):
        if self.filename_box.GetValue().endswith(".scans"):
            self.hdist_txtbox.SetValue("")
            self.hdist_err_txtbox.SetValue("")
            self.hdist_txtbox.Disable()
            self.hdist_err_txtbox.Disable()
            self.vangle_txtbox.SetValue("")
            self.vangle_txtbox.Disable()
            self.vangle_err_txtbox.SetValue("")
            self.vangle_err_txtbox.Disable()
            self.conv_txtbox.SetValue("")
            self.conv_txtbox.Disable()
            self.method_choice.Disable()
            
        else:
            self.hdist_txtbox.Enable()
            self.hdist_err_txtbox.SetValue("0.0")
            self.hdist_err_txtbox.Enable()
            self.vangle_txtbox.Enable()
            self.vangle_err_txtbox.SetValue("0.0")
            self.vangle_err_txtbox.Enable()    
            self.conv_txtbox.SetValue("1.0")
            self.conv_txtbox.Enable()
            self.method_choice.Enable()
            
    def on_ok(self, evnt):
        filename = self.filename_box.GetValue()
        
        if not os.path.exists(filename):
            wx.MessageBox("File does not exist!", "ICA Inspect", wx.ICON_ERROR)
            return
             
        wx.YieldIfNeeded()
        if filename.endswith(".scans"):
            self.Show(False)
            with open(filename, "rb") as ifp:
                scans = cPickle.load(ifp)

        else:
            try:
                h_dist = float(self.hdist_txtbox.GetValue())
            except ValueError:
                wx.MessageBox("Horizontal distance must be a float (meters).", "ICA Inspect", wx.ICON_ERROR)
                return
            
            try:
                h_dist_err = float(self.hdist_err_txtbox.GetValue())
            except ValueError:
                wx.MessageBox("Horizontal distance error must be a float (meters).", "ICA Inspect", wx.ICON_ERROR)
                return
            
            try:
                v_angle = float(self.vangle_txtbox.GetValue())
            except ValueError:
                wx.MessageBox("Vertical angle must be a float (degrees).", "ICA Inspect", wx.ICON_ERROR)
                return
            
            try:
                v_angle_err = float(self.vangle_err_txtbox.GetValue())
            except ValueError:
                wx.MessageBox("Vertical angle error must be a float (degrees).", "ICA Inspect", wx.ICON_ERROR)
                return
            
            try:
                conv = float(self.conv_txtbox.GetValue())
            except ValueError:
                wx.MessageBox("Conversion factor must be a float.", "ICA Inspect", wx.ICON_ERROR)
                return
            
            int_method = self.method_choice.GetSelection()
            
            self.Show(False)
            scans = []
            #Note that units conversion is done here
            for i in fullscans(filename, conv=conv):
                #scans need to have at least 4 points in them - otherwise we can't fit a gaussian to them
                if len(i["angle"]) < 4:
                    continue
                scans.append(ProcessedScan(i["angle"].copy(), i["times"].copy(),i["so2_col_amount"].copy(), i["angle_error"].copy(), i["so2_col_amount_error"].copy(), i["int_times"].copy(), i["coadds"].copy(), h_dist, h_dist_err, v_angle, v_angle_err, self.int_methods[int_method]))
            print "Using int method:",self.int_methods[int_method]
            #do the fittings/integrations
            scans = parallel_process(calc_scan, scans)
        
        self.p = PlotManager(scans)    
        ctrl_win = ControlWindow(self.p, os.path.basename(filename)+" - ICA Inspect")
        self.p.show()
        self.Destroy()

def calc_scan(x):
    return x.calculate() 
 